//
// Created by limo on 2025/9/27.
//

#pragma once

#include <iostream>
#include <string>
#include <map>
#include <memory>
#include <chrono>

#include <boost/asio.hpp>
#include <boost/bind/bind.hpp>

namespace net {
    namespace asio = boost::asio;
    using boost::asio::ip::udp;

    struct UdpSession {
        std::string clientAddress;
        uint16_t clientPort;
        uint32_t sessionId;
        std::chrono::steady_clock::time_point lastActiveTime;
        boost::asio::ip::udp::endpoint clientEndpoint;
    };

    // --- 回调函数类型定义 ---

    // 1. 连接回调：void(const UdpSession&)
    using ConnectCallback = std::function<void(const UdpSession&)>;
    // 2. 断线回调：void(const UdpSession&)
    using DisconnectCallback = std::function<void(const UdpSession&)>;
    // 3. 消息回调：void(UdpSession&, const std::string&, std::size_t)
    using UdpMessageCallback = std::function<void(UdpSession&, const std::string&, std::size_t)>;
    // 4. 错误回调：void(const boost::system::error_code&)
    using ErrorCallback = std::function<void(const boost::system::error_code&)>;

    class UdpServer {
    public:
        // 构造函数：现在只需要 io_context 和端口
        UdpServer(asio::io_context& ioContext, short listenPort)
                : m_socket(ioContext, udp::endpoint(udp::v4(), listenPort)),
                  m_sessionCleanerTimer(ioContext),
                  m_ioContext(ioContext),
                  m_nextSessionId(1)
        {
            std::cout << "udpServer started on port " << listenPort << std::endl;
            startReceive();
            startCleanupTimer();
        }

        // --- 公有方法：设置回调函数 ---

        void setConnectCallback(ConnectCallback cb)    { m_onConnect = std::move(cb); }
        void setDisconnectCallback(DisconnectCallback cb) { m_onDisconnect = std::move(cb); }
        void setMessageCallback(UdpMessageCallback cb)    { m_onMessage = std::move(cb); }
        void setErrorCallback(ErrorCallback cb)      { m_onError = std::move(cb); }

        // 公有方法：用于应用逻辑向客户端发送数据
        void sendTo(const std::string& message, const udp::endpoint& targetEndpoint) {
            // ... (实现与前文相同，保证异步发送)
            auto sharedMessage = std::make_shared<std::string>(message);

            m_socket.async_send_to(
                    asio::buffer(*sharedMessage),
                    targetEndpoint,
                    boost::bind(&UdpServer::handleSend, this,
                                sharedMessage,
                                asio::placeholders::error,
                                asio::placeholders::bytes_transferred)
            );
        }

    private:
        // 核心成员变量
        udp::socket m_socket;
        udp::endpoint m_remoteEndpoint;
        std::array<char, 1024> m_recvBuffer;
        std::map<udp::endpoint, UdpSession> m_activeSessions;
        asio::steady_timer m_sessionCleanerTimer;
        asio::io_context& m_ioContext;
        uint32_t m_nextSessionId;

        static constexpr std::chrono::seconds SESSION_TIMEOUT = std::chrono::seconds(10);

        // --- 核心改动：回调成员变量 ---
        ConnectCallback    m_onConnect;
        DisconnectCallback m_onDisconnect;
        UdpMessageCallback    m_onMessage;
        ErrorCallback      m_onError;

        // --- 私有实现函数 (仅包含触发回调的逻辑) ---

        void startReceive() {
            m_socket.async_receive_from(
                    asio::buffer(m_recvBuffer), m_remoteEndpoint,
                    boost::bind(&UdpServer::handleReceive, this, asio::placeholders::error, asio::placeholders::bytes_transferred)
            );
        }

        void handleReceive(const boost::system::error_code& error, std::size_t bytesTransferred) {
            if (!error) {
                auto it = m_activeSessions.find(m_remoteEndpoint);
                UdpSession* sessionPtr = nullptr;
                bool isNewSession = (it == m_activeSessions.end());

                if (isNewSession) {
                    // 1. 新会话：创建并初始化
                    UdpSession newSession;
                    newSession.clientAddress = m_remoteEndpoint.address().to_string();
                    newSession.clientPort = m_remoteEndpoint.port();
                    newSession.sessionId = m_nextSessionId++;
                    newSession.clientEndpoint = m_remoteEndpoint;
                    newSession.lastActiveTime = std::chrono::steady_clock::now();

                    it = m_activeSessions.emplace(m_remoteEndpoint, std::move(newSession)).first;
                    sessionPtr = &it->second;

                    // 触发 onConnect
                    if (m_onConnect) m_onConnect(*sessionPtr);
                } else {
                    // 现有会话：更新时间
                    sessionPtr = &it->second;
                    sessionPtr->lastActiveTime = std::chrono::steady_clock::now();
                }

                // 2. 触发 onMessage
                std::string message(m_recvBuffer.data(), bytesTransferred);
                if (m_onMessage) m_onMessage(*sessionPtr, message, bytesTransferred);

            } else {
                // 3. 触发 onError
                if (m_onError) m_onError(error);
            }

            startReceive();
        }

        void handleSend(std::shared_ptr<std::string> /*message*/,
                        const boost::system::error_code& error,
                        std::size_t /*bytesTransferred*/) {
            if (error) {
                std::cerr << "[ERROR] Send error: " << error.message() << std::endl;
                if (m_onError) m_onError(error);
            }
        }

        void startCleanupTimer() {
            m_sessionCleanerTimer.expires_after(std::chrono::seconds(5));
            m_sessionCleanerTimer.async_wait(
                    boost::bind(&UdpServer::handleCleanupTimer, this, asio::placeholders::error)
            );
        }

        void handleCleanupTimer(const boost::system::error_code& error) {
            if (!error) {
                auto now = std::chrono::steady_clock::now();

                for (auto it = m_activeSessions.begin(); it != m_activeSessions.end(); ) {
                    if (now - it->second.lastActiveTime > SESSION_TIMEOUT) {
                        // 触发 onDisconnect
                        if (m_onDisconnect) m_onDisconnect(it->second);

                        it = m_activeSessions.erase(it);
                    } else {
                        ++it;
                    }
                }
                startCleanupTimer();
            }
        }
    };
} // net
